"""
层级结构定义模块
"""

from typing import Dict, List, Tuple, Optional
import pandas as pd
import numpy as np


class HierarchicalStructure:
    """层级结构管理类"""
    
    def __init__(self, config: Dict = None):
        """
        初始化层级结构
        
        Args:
            config: 配置字典
        """
        self.config = config or {}
        self.regions = self.config.get('hierarchy', {}).get('regions', {})
        self.levels = self.config.get('hierarchy', {}).get('levels', [])
        
    def get_region_provinces(self, region: str) -> List[str]:
        """
        获取指定区域的省份列表
        
        Args:
            region: 区域名称
            
        Returns:
            省份列表
        """
        return self.regions.get(region, [])
    
    def get_province_region(self, province: str) -> Optional[str]:
        """
        获取省份所属区域
        
        Args:
            province: 省份名称
            
        Returns:
            区域名称
        """
        for region, provinces in self.regions.items():
            if province in provinces:
                return region
        return None
    
    def get_all_provinces(self) -> List[str]:
        """
        获取所有省份列表
        
        Returns:
            所有省份列表
        """
        all_provinces = []
        for provinces in self.regions.values():
            all_provinces.extend(provinces)
        return all_provinces
    
    def get_all_regions(self) -> List[str]:
        """
        获取所有区域列表
        
        Returns:
            所有区域列表
        """
        return list(self.regions.keys())
    
    def create_aggregation_matrix(self) -> np.ndarray:
        """
        创建聚合矩阵
        
        Returns:
            聚合矩阵
        """
        # 获取所有省份和区域
        all_provinces = self.get_all_provinces()
        all_regions = self.get_all_regions()
        
        # 创建聚合矩阵
        n_provinces = len(all_provinces)
        n_regions = len(all_regions)
        n_total = n_provinces + n_regions + 1  # +1 for national level
        
        # 初始化矩阵
        S = np.zeros((n_total, n_total))
        
        # 省份到区域的映射
        province_to_region = {}
        for region, provinces in self.regions.items():
            for province in provinces:
                province_to_region[province] = region
        
        # 填充矩阵
        # 省份索引
        province_indices = {province: i for i, province in enumerate(all_provinces)}
        # 区域索引
        region_indices = {region: i + n_provinces for i, region in enumerate(all_regions)}
        # 全国索引
        national_index = n_provinces + n_regions
        
        # 省份到区域
        for province, region in province_to_region.items():
            if province in province_indices and region in region_indices:
                p_idx = province_indices[province]
                r_idx = region_indices[region]
                S[p_idx, r_idx] = 1
        
        # 区域到全国
        for region in all_regions:
            if region in region_indices:
                r_idx = region_indices[region]
                S[r_idx, national_index] = 1
        
        return S
    
    def validate_hierarchy(self, data: pd.DataFrame) -> bool:
        """
        验证层级结构的一致性
        
        Args:
            data: 数据框
            
        Returns:
            是否一致
        """
        # 检查省份是否在定义的区域中
        if 'Province' in data.columns and 'Area' in data.columns:
            for _, row in data.iterrows():
                province = row['Province']
                area = row['Area']
                
                if area != '全国':  # 跳过全国级别
                    expected_area = self.get_province_region(province)
                    if expected_area and expected_area != area:
                        return False
        
        return True
    
    def get_level_data(self, data: pd.DataFrame, level: str) -> pd.DataFrame:
        """
        获取指定层级的数据
        
        Args:
            data: 原始数据
            level: 层级名称 ('province', 'region', 'national')
            
        Returns:
            指定层级的数据
        """
        if level == 'province':
            return data[data['Province'] != data['Area']]
        elif level == 'region':
            return data[data['Province'] == data['Area']]
        elif level == 'national':
            return data[data['Area'] == '全国']
        else:
            raise ValueError(f"不支持的层级: {level}")
    
    def aggregate_data(self, data: pd.DataFrame, target_level: str) -> pd.DataFrame:
        """
        聚合数据到指定层级
        
        Args:
            data: 原始数据
            target_level: 目标层级
            
        Returns:
            聚合后的数据
        """
        if target_level == 'region':
            # 聚合到区域级别
            aggregated = data.groupby(['Year', 'Area'])['Grain'].sum().reset_index()
            aggregated['Province'] = aggregated['Area']
            return aggregated
        elif target_level == 'national':
            # 聚合到全国级别
            aggregated = data.groupby('Year')['Grain'].sum().reset_index()
            aggregated['Area'] = '全国'
            aggregated['Province'] = '全国'
            return aggregated
        else:
            raise ValueError(f"不支持的目标层级: {target_level}")
    
    def get_hierarchy_info(self) -> Dict:
        """
        获取层级结构信息
        
        Returns:
            层级结构信息字典
        """
        info = {
            'regions': self.regions,
            'levels': self.levels,
            'total_provinces': len(self.get_all_provinces()),
            'total_regions': len(self.get_all_regions()),
            'hierarchy_depth': len(self.levels)
        }
        return info 